import os
import numpy as np
import librosa
import warnings
warnings.filterwarnings('ignore')

def convert_to_stereo_16bit(audio_data, target_sample_rate=22050):
    """
    将音频数据转换为双通道22.05K 16bits格式
    参数:
        audio_data: 单声道音频数据
        target_sample_rate: 目标采样率
    返回:
        16位整型的双通道音频数据
    """
    # 确保音频数据是单声道
    if len(audio_data.shape) > 1:
        # 如果是多声道，取第一个声道
        audio_data = audio_data[:, 0]
    
    # 将单声道复制为双声道（左右声道相同）
    stereo_data = np.column_stack((audio_data, audio_data))
    
    # 转换为16位整型（范围：-32768 到 32767）
    # 先限制在[-1, 1]范围内，然后缩放到16位整型范围
    stereo_data = np.clip(stereo_data, -1.0, 1.0)
    stereo_data_16bit = (stereo_data * 32767).astype(np.int16)
    
    return stereo_data_16bit

def convert_mp3_to_bin_stereo_16bit():
    """
    使用librosa库转换MP3到双通道22.05K 16bits的BIN文件
    """
    input_folder = "按键音"
    output_folder = "原始音频数据"
    target_sample_rate = 22050
    
    # 检查输入文件夹
    if not os.path.exists(input_folder):
        print(f"错误: 找不到 '{input_folder}' 文件夹")
        return
    
    # 创建输出文件夹
    os.makedirs(output_folder, exist_ok=True)
    
    # 获取所有MP3文件
    mp3_files = [f for f in os.listdir(input_folder) if f.lower().endswith('.mp3')]
    
    if not mp3_files:
        print(f"在 '{input_folder}' 文件夹中没有找到MP3文件")
        return
    
    print(f"找到 {len(mp3_files)} 个MP3文件，开始转换...")
    print(f"目标格式: 双通道, {target_sample_rate}Hz, 16bits")
    print("使用librosa处理...")
    
    success_count = 0
    failed_files = []
    total_file_size = 0
    
    for mp3_file in mp3_files:
        try:
            input_path = os.path.join(input_folder, mp3_file)
            output_filename = os.path.splitext(mp3_file)[0] + '.bin'
            output_path = os.path.join(output_folder, output_filename)
            
            print(f"正在处理: {mp3_file}", end="")
            
            # 使用librosa加载音频文件，直接重采样到目标采样率
            audio_data, _ = librosa.load(
                input_path, 
                sr=target_sample_rate,  # 直接重采样到目标采样率
                mono=True,              # 强制单声道（后面会转为双声道）
                res_type='kaiser_fast'  # 快速重采样
            )
            
            # 转换为双通道16bits格式
            stereo_16bit_data = convert_to_stereo_16bit(audio_data, target_sample_rate)
            
            # 保存为二进制文件（16位整型）
            with open(output_path, 'wb') as bin_file:
                stereo_16bit_data.tofile(bin_file)
            
            # 计算文件信息
            file_size = os.path.getsize(output_path)
            total_file_size += file_size
            duration = len(audio_data) / target_sample_rate
            file_size_kb = file_size / 1024
            
            print(f" -> ✓ 完成")
            print(f"     时长: {duration:.2f}秒, 样本数: {len(stereo_16bit_data)}")
            print(f"     格式: 双通道 16bits, 文件大小: {file_size_kb:.1f}KB")
            print(f"     输出: {output_filename}")
            
            success_count += 1
            
        except Exception as e:
            print(f" -> ✗ 失败: {str(e)}")
            failed_files.append(mp3_file)
    
    print(f"\n转换完成!")
    print(f"成功: {success_count}/{len(mp3_files)}")
    
    if success_count > 0:
        total_size_mb = total_file_size / (1024 * 1024)
        avg_size_kb = total_file_size / success_count / 1024
        print(f"总输出大小: {total_size_mb:.2f}MB")
        print(f"平均文件大小: {avg_size_kb:.1f}KB")
    
    if failed_files:
        print(f"失败: {len(failed_files)}")
        print("\n失败的文件:")
        for f in failed_files:
            print(f"  - {f}")
    
    print(f"\n输出文件保存在: {output_folder}")

def verify_bin_file(file_path):
    """
    验证BIN文件格式是否正确
    """
    try:
        # 读取二进制文件
        with open(file_path, 'rb') as f:
            data = np.fromfile(f, dtype=np.int16)
        
        # 检查数据长度是否为偶数（双通道）
        if len(data) % 2 != 0:
            print(f"错误: 文件 {file_path} 数据长度不是偶数")
            return False
        
        # 重塑为双通道
        stereo_data = data.reshape(-1, 2)
        
        print(f"文件: {os.path.basename(file_path)}")
        print(f"  总样本数: {len(data)}")
        print(f"  帧数: {len(stereo_data)}")
        print(f"  声道数: 2")
        print(f"  数据类型: int16")
        print(f"  数值范围: {np.min(data)} 到 {np.max(data)}")
        
        return True
        
    except Exception as e:
        print(f"验证文件 {file_path} 时出错: {e}")
        return False

# 测试转换函数
def test_conversion():
    """测试音频转换功能"""
    print("测试音频格式转换...")
    
    # 创建测试数据：1秒的440Hz正弦波
    sample_rate = 22050
    duration = 1.0
    t = np.linspace(0, duration, int(sample_rate * duration), endpoint=False)
    test_audio = 0.5 * np.sin(2 * np.pi * 440 * t)  # 440Hz正弦波
    
    print(f"原始数据: 单声道, {len(test_audio)} 样本, float32")
    
    # 转换为双通道16bits
    stereo_16bit = convert_to_stereo_16bit(test_audio, sample_rate)
    
    print(f"转换后: 双通道, {len(stereo_16bit)} 样本, int16")
    print(f"  形状: {stereo_16bit.shape}")
    print(f"  左声道范围: {np.min(stereo_16bit[:, 0])} 到 {np.max(stereo_16bit[:, 0])}")
    print(f"  右声道范围: {np.min(stereo_16bit[:, 1])} 到 {np.max(stereo_16bit[:, 1])}")

if __name__ == "__main__":
    # 可以选择先运行测试
    # test_conversion()
    # print("\n" + "="*50 + "\n")
    
    convert_mp3_to_bin_stereo_16bit()
    
    # 转换完成后可以选择验证一个文件
    # print("\n" + "="*50)
    # print("验证输出文件格式...")
    # output_folder = "原始音频数据"
    # if os.path.exists(output_folder):
    #     bin_files = [f for f in os.listdir(output_folder) if f.endswith('.bin')]
    #     if bin_files:
    #         verify_bin_file(os.path.join(output_folder, bin_files[0]))
